{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61PmYiKzWgz4"
      },
      "source": [
        "# Illustrate how to use a PyTorch Gemma model running locally.\n",
        "\n",
        "Note: This will download a 2 billion parameter model from Hugging Face, so make sure you have enough space for that."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "evAcCqotVPWY"
      },
      "outputs": [],
      "source": [
        "# @title Colab-specific setup (use a CodeSpace to avoid the need for this).\n",
        "try:\n",
        "  %env COLAB_RELEASE_TAG\n",
        "except:\n",
        "  pass  # Not running in colab.\n",
        "else:\n",
        "  %pip install --ignore-requires-python --requirement 'https://raw.githubusercontent.com/google-deepmind/concordia/main/examples/requirements.in' 'git+https://github.com/google-deepmind/concordia.git#egg=gdm-concordia'\n",
        "  %pip list"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AVD1XXzoU5-o"
      },
      "outputs": [],
      "source": [
        "from concordia.contrib.language_models.huggingface import pytorch_gemma_model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e28CPlxZViAQ"
      },
      "outputs": [],
      "source": [
        "# This will download the model from Hugging Face.\n",
        "model = pytorch_gemma_model.PyTorchGemmaLanguageModel(\n",
        "    model_name='google/gemma-2b-it',\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DT8QPSHDV36C"
      },
      "outputs": [],
      "source": [
        "choice_response = model.sample_choice('What comes next? a,b,c,d,e,f,',\n",
        "                                      responses=['d', 'g', 'z'])\n",
        "choice_response[1]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qDKp1LeWWUfZ"
      },
      "outputs": [],
      "source": [
        "text_response = model.sample_text('What is the meaning of life?',\n",
        "                                  max_tokens=40)\n",
        "text_response"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RMj2IZffYYh7"
      },
      "source": [
        "```\n",
        "Copyright 2023 DeepMind Technologies Limited.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "    https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License.\n",
        "```"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
